from main.german_transliterate.core import GermanTransliterate
from paddlespeech.t2s.frontend.phonectic import English
from paddlespeech.t2s.frontend.punctuation import get_punctuations
from paddlespeech.t2s.frontend.vocab import Vocab


class GermanyG2P:
    punctuations = get_punctuations(lang='en')
    phonemes = "a aj aw aː b c cʰ d eː f h iː j k kʰ l l̩ m m̩ n n̩ oː p pf pʰ s t ts tʃ tʰ uː v x yː z ç øː ŋ œ ɐ ɔ " \
               "ɔʏ ə ɛ ɟ ɡ ɪ ɲ ʁ ʃ ʊ ʏ".split(' ')

    def __init__(self, germany_dict_path, phone_id_map_path):
        self.dict = {}
        self.normalizer = GermanTransliterate(replace={';': ',', ':': ' '}, sep_abbreviation=' -- ')
        with open(germany_dict_path, 'r') as f:
            for l in f.readlines():
                l = l[:-1]
                seg_l = l.split('\t')
                self.dict[seg_l[0]] = seg_l[-1]
        self.phone_id_map = {}
        with open(phone_id_map_path, 'r') as f:
            for l in f.readlines():
                l = l[:-1]
                seg_l = l.split(' ')
                self.phone_id_map[seg_l[0]] = int(seg_l[1])

    def __call__(self, sentence: str):
        sentence = self.normalizer.transliterate(sentence)
        phones = []
        words = []
        word = ''
        for s in sentence:
            if s in self.punctuations:
                if len(word) > 0:
                    words.append(word)
                    word = ''
                    if s != ' ':
                        words.append(s)
                continue
            word += s

        if len(word) > 0:
            words.append(word)

        for word in words:

            phone = self.dict.get(word, None)
            if phone is None:
                if word in self.punctuations:
                    phones.extend(['sp'])
                else:
                    phones.extend(['spn'])
            else:
                phones.extend(phone.split(' '))
        return phones


class Germany(English):

    def __init__(self, germany_dict_path, phone_vocab_path):
        super().__init__(phone_vocab_path)
        self.backend = GermanyG2P(germany_dict_path, phone_vocab_path)
        self.phonemes = list(self.backend.phonemes)
        self.punctuations = self.backend.punctuations
        self.vocab = Vocab(self.phonemes + self.punctuations)


if __name__ == '__main__':
    #     g2p = GermanyG2P(germany_dict_path='/home/german_mfa.dict', phone_id_map_path='/home/output/dump1/phone_id_map.txt')
    #     # phones = g2p(['feuer'])
    #     # print(phones)

    # result = normalizer.transliterate('Okay, sicher. Hier ist eine große Auswahl aus meiner Bibliothek.')
    #     phones = g2p(result)
    #     input_id = [g2p.phone_id_map[p] for p in phones]
    front_end = Germany(germany_dict_path='/home/dicts/german_mfa.dict',
                        phone_vocab_path='/home/output/dump1/phone_id_map.txt')
    print(front_end.punctuations)
    result = front_end.get_input_ids(sentence='Okay, sicher. Hier ist eine große Auswahl aus meiner Bibliothek.',
                                     merge_sentences=True)
    print(result)
